package com.chinamobile.iot.handler.impl;

import com.alibaba.fastjson.JSON;
import com.chinamobile.iot.handler.exception.ZkException;
import com.chinamobile.iot.handler.model.BrokerInfo;
import com.chinamobile.iot.handler.model.BrokerStat;
import com.chinamobile.iot.handler.model.ZkClientInfo;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.CountDownLatch;

import static com.chinamobile.iot.handler.model.BrokerInfo.ABNORMAL;
import static com.chinamobile.iot.handler.model.BrokerInfo.NORMAL;
import static org.apache.zookeeper.Watcher.Event.EventType.NodeChildrenChanged;

/**
 * Created by szl on 2017/7/11.
 */
public class ZkHelper {

    private final static String ZNODE_BROKERS = "/brokers";
    private final static String ZNODE_CLIENTS = "/clients";
    private final static String ZNODE_MANAGER = "/manager";
    private final static String ZNODE_MANAGER_FULL = "/manager/actor";
    private final static String ZNODE_MANAGER_DATA = "/manager/data";

    private final static Logger logger = LoggerFactory.getLogger(ZkHelper.class);
    private final BrokerWatcher brokerWatcher = new BrokerWatcher(this);
    private final ZkWatcher zkWatcher = new ZkWatcher(this);
    private final BrokerMonitor brokerMonitor;
    private final String connString;
    private final ZkConnectEventHandler zkConnectEventHandler;
    private final String brokerName;
    private ZooKeeper zooKeeper;
    private CountDownLatch connectOk;
    private String managerBrokerName;


    /**
     * init the zkHelper
     * @param connString the zookeeper server address
     * @param monitor the monitor on znode /brokers, if new broker is add of some broker is disconnected, it will be callback
     * @param connectEventHandler the EventHandler when the connection to the zookeeper is just connected
     */
    public ZkHelper(String connString, BrokerMonitor monitor, ZkConnectEventHandler connectEventHandler, String brokerName) {
        this.brokerMonitor = monitor;
        this.connString = connString;
        this.zkConnectEventHandler = connectEventHandler;
        this.brokerName = brokerName;
    }

    public void init() throws ZkException {
        try {
            connectOk = new CountDownLatch(1);
            zooKeeper = new ZooKeeper(connString, 30000, zkWatcher);
            connectOk.await();
            Stat exists = zooKeeper.exists(ZNODE_BROKERS, false);
            if (null == exists) {
                createRoot(ZNODE_BROKERS);
            }

            exists = zooKeeper.exists(ZNODE_CLIENTS, false);
            if (null == exists) {
                createRoot(ZNODE_CLIENTS);
            }
            exists = zooKeeper.exists(ZNODE_MANAGER, false);
            if (null == exists) {
                createRoot(ZNODE_MANAGER);
            }
            brokerMonitor.setZkHelper(this);
            // watch all child of broker
            zooKeeper.getChildren(ZNODE_BROKERS, brokerWatcher);
        } catch (Exception e) {
            logger.error("ZkHelper::onConnected(...) failed, {}", e.getMessage());
        }

        boolean b = zkConnectEventHandler.onConnected(this);
        if (!b) {
            throw new ZkException("the broker is maintaining, please startup later...");
        }
    }

    public void destroy() {
        if (zooKeeper != null) {
            try {
                zooKeeper.close();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    public ZooKeeper getZooKeeper() {
        return zooKeeper;
    }

    public boolean isConnectedToZookeeper() {
        ZooKeeper.States states = zooKeeper.getState();
        return states.isConnected();
    }

    private boolean createRoot(String path) {
        boolean result = false;
        try {
            zooKeeper.create(path, null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            result = true;
        } catch (Exception e) {
            return false;
        }
        return result;
    }

    public boolean addBroker(String name, BrokerInfo brokerInfo) {
        StringBuilder stringBuilder = new StringBuilder(128);
        stringBuilder.append(ZNODE_BROKERS);
        stringBuilder.append("/");
        stringBuilder.append(name);

        String path = stringBuilder.toString();
        byte[] data = JSON.toJSONString(brokerInfo).getBytes();

        try {
            zooKeeper.create(path, data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
            addBrokerToManagerData(name);
        } catch (Exception e) {
            logger.warn("add broker {} failed, brokerinfo: {}", JSON.toJSONString(brokerInfo), e.getMessage());
            return false;
        }
        return true;
    }

    private void addBrokerToManagerData(String name){
        try{
            Stat stat = zooKeeper.exists(ZNODE_MANAGER_DATA, false);
            if (stat == null){
                try{
                    zooKeeper.create(ZNODE_MANAGER_DATA, name.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
                    return;
                }
                catch (Exception e){
                    logger.info("the zk path {} already exists.", ZNODE_MANAGER_DATA);
                }
            }
            logger.info("begin modify the data");
            while (true) {
                try {
                    Stat stat1 = new Stat();
                    byte[] dd = zooKeeper.getData(ZNODE_MANAGER_DATA, false, stat1);
                    String strData = new String(dd);
                    String[] ary = strData.split(";");
                    boolean exist = false;
                    for (int i = 0; i<ary.length; i++){
                        if (ary[i].equals(name)){
                            exist = true;
                            break;
                        }
                    }
                    if (!exist){
                        strData = strData + ";" + name;
                        zooKeeper.setData(ZNODE_MANAGER_DATA, strData.getBytes(), stat1.getVersion());
                        break;
                    }
                    else{
                        break;
                    }
                } catch (Exception e) {
                }
            }
        }
        catch (Exception e){

        }
    }


    public BrokerInfo getBroker(String name) {
        String path = ZNODE_BROKERS + "/" + name;
        try {
            Stat stat = new Stat();
            byte[] data = zooKeeper.getData(path, false, stat);
            if (data != null) {
                String value = new String(data);
                BrokerInfo info = JSON.parseObject(value, BrokerInfo.class);
                return info;
            }
            return null;
        } catch (Exception e) {
            return null;
        }
    }

    public boolean isBrokerValid(String name) {
        String path = ZNODE_BROKERS + "/" + name;

        try {
            Stat stat = zooKeeper.exists(path, false);
            if (stat != null) {
                byte[] data = zooKeeper.getData(path, false, stat);
                if (data != null) {
                    String value = new String(data);
                    BrokerInfo info = JSON.parseObject(value, BrokerInfo.class);
                    if (info.getStatus().equals(NORMAL)) {
                        return true;
                    } else {
                        return false;
                    }
                }
                return false;
            }
        } catch (Exception e) {
            return false;
        }
        return false;
    }

    public BrokerStat getBrokerStat(String name) {
        String path = ZNODE_BROKERS + "/" + name;

        try {
            Stat stat = zooKeeper.exists(path, false);
            if (stat != null) {
                byte[] data = zooKeeper.getData(path, false, stat);
                if (data != null) {
                    String value = new String(data);
                    BrokerInfo info = JSON.parseObject(value, BrokerInfo.class);
                    if (info.getStatus().equals(NORMAL)) {
                        if (stat.getEphemeralOwner() == zooKeeper.getSessionId()) {
                            return BrokerStat.NORMAL;
                        } else {
                            return BrokerStat.OLD;
                        }
                    } else {
                        return BrokerStat.MAINTAIN;
                    }
                } else {
                    return BrokerStat.NOEXIST;
                }
            } else {
                return BrokerStat.NOEXIST;
            }
        } catch (Exception e) {
            return BrokerStat.NOEXIST;
        }
    }

    public long getSessionId() {
        return zooKeeper.getSessionId();
    }


    public boolean beginCleanBroker(String name) {
        StringBuilder stringBuilder = new StringBuilder(128);
        stringBuilder.append(ZNODE_BROKERS);
        stringBuilder.append("/");
        stringBuilder.append(name);
        String path = stringBuilder.toString();

        BrokerInfo brokerInfo = new BrokerInfo();
        brokerInfo.setName(name);
        brokerInfo.setStatus(ABNORMAL);
        byte[] data = JSON.toJSONString(brokerInfo).getBytes();
        try {
            zooKeeper.create(path, data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            return true;
        } catch (Exception e) {
            logger.error("create node {} to maintain failed. {}", path, e.getMessage());
        }
        return false;
    }




    public void finishCleanBroker(String name) {
        StringBuilder stringBuilder = new StringBuilder(128);
        stringBuilder.append(ZNODE_BROKERS);
        stringBuilder.append("/");
        stringBuilder.append(name);
        String path = stringBuilder.toString();
        try {
            Stat stat = zooKeeper.exists(path, false);
            if (stat != null) {
                byte[] data = zooKeeper.getData(path, null, stat);
                if (data != null) {
                    String value = new String(data);
                    BrokerInfo brokerInfo = JSON.parseObject(value, BrokerInfo.class);
                    if (!brokerInfo.getStatus().equals(NORMAL)) {
                        zooKeeper.delete(path, stat.getVersion());
                    }
                }
            }
        } catch (Exception e) {
            logger.error("finishCleanBroker failed. remove {} failed, {}", path, e.getMessage());
        }
    }

    public void removeBroker(String name) {
        StringBuilder stringBuilder = new StringBuilder(128);
        stringBuilder.append(ZNODE_BROKERS);
        stringBuilder.append("/");
        stringBuilder.append(name);
        String path = stringBuilder.toString();
        try {
            Stat stat = zooKeeper.exists(path, false);
            if (stat != null) {
                zooKeeper.delete(path, stat.getVersion());
            }
        } catch (Exception e) {
            logger.error("finishCleanBroker failed. remove {} failed, {}", path, e.getMessage());
        }
    }

    public boolean addClient(String mqttId, ZkClientInfo clientInfo) {
        StringBuilder stringBuilder = new StringBuilder(128);
        stringBuilder.append(ZNODE_CLIENTS);
        stringBuilder.append("/");
        stringBuilder.append(mqttId);
        String path = stringBuilder.toString();

        String value = JSON.toJSONString(clientInfo);
        byte[] data = value.getBytes();

        boolean result = false;

        try {
            zooKeeper.create(path, data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
            result = true;
        } catch (KeeperException e) {
            logger.error("addClient failed. the client alreay exists, {}, {}", mqttId, e.getMessage());
            result = false;
        } catch (Exception e) {
            logger.error("addClient failed. {}", mqttId, e.getMessage());
            result = false;
        }
        return result;
    }


    public boolean removeClient(String mqttId) {
        logger.info("ZkHelper::removeClient is called. mqttId = {}", mqttId);
        StringBuilder stringBuilder = new StringBuilder(128);
        stringBuilder.append(ZNODE_CLIENTS);
        stringBuilder.append("/");
        stringBuilder.append(mqttId);
        String path = stringBuilder.toString();

        boolean result = false;

        try {
            Stat stat = zooKeeper.exists(path, false);
            if (stat != null) {
                zooKeeper.delete(path, stat.getVersion());
            }
            result = true;
        } catch (Exception e) {
            logger.error("removeClient error, clientId:{}, cause:{}", mqttId, e.getMessage());
        }
        return result;
    }

    public ZkClientInfo getClient(String mqttId) {
        StringBuilder stringBuilder = new StringBuilder(128);
        stringBuilder.append(ZNODE_CLIENTS);
        stringBuilder.append("/");
        stringBuilder.append(mqttId);
        String path = stringBuilder.toString();
        try {
            byte[] data = zooKeeper.getData(path, false, null);
            if (data != null) {
                String value = new String(data);
                ZkClientInfo info = JSON.parseObject(value, ZkClientInfo.class);
                return info;
            }
            return null;
        } catch (Exception e) {
            return null;
        }
    }

    public boolean tryBecomeManager(String name) {
        String path = ZNODE_MANAGER_FULL;
        if (name == null){
            name = brokerName;
        }
        try {
            zooKeeper.create(path, name.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
            return true;
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * remove the manager if the manager name is name
     * @param name
     */
    public void removeManager(String name){
        try {
            Stat stat = zooKeeper.exists(ZNODE_MANAGER_FULL, false);
            if (stat != null){
                byte[] data = zooKeeper.getData(ZNODE_MANAGER_FULL, false, stat);
                if (data != null){
                    String broker = new String(data);
                    if (broker.equals(name)){
                        zooKeeper.delete(ZNODE_MANAGER_FULL, stat.getVersion());
                    }
                }
            }
        }
        catch (Exception e){
        }
    }

    public void setManagerData(String data){
        String path = ZNODE_MANAGER_DATA;
        try {
            Stat stat = zooKeeper.exists(path, false);
            if (stat == null){
                zooKeeper.create(path, data.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
            else{
                zooKeeper.setData(path, data.getBytes(), stat.getVersion());
            }
        } catch (Exception e) {
        }
    }

    public String getManagerData(){
        String path = ZNODE_MANAGER_DATA;
        try{
            byte[] data = zooKeeper.getData(path, false, null);
            if (data != null){
                return new String(data);
            }
        }
        catch (Exception e){

        }
        return null;
    }



    public String getManager(){
        String path = ZNODE_MANAGER_FULL;
        try {
            byte[] data = zooKeeper.getData(path, false, null);
            if (data != null){
                return new String(data);
            }
        }
        catch (Exception e){
        }
        return null;
    }

    public interface ZkConnectEventHandler {
        boolean onConnected(ZkHelper zkHelper);
    }

    private static class ZkWatcher implements Watcher {

        private ZkHelper zkHelper;
        private int retryCount = 0;

        public ZkWatcher(ZkHelper zkHelper) {
            this.zkHelper = zkHelper;
        }

        @Override
        public void process(WatchedEvent event) {
            switch (event.getState()) {
                case SyncConnected: {
                    zkHelper.connectOk.countDown();
                }
                break;
                case Disconnected: {
                    logger.info("the connection to zookeeper is disconnected!");
                }
                break;
                case AuthFailed: {
                    logger.info("the connection to zookeeper auth failed.");
                }
                break;
                case ConnectedReadOnly: {
                    logger.info("the connection to zookeeper with readonly mode");
                }
                break;
                case SaslAuthenticated: {
                    logger.info("the connection to zookeeper with saslAuthenticated!");
                }
                break;
                case Expired: {
                    // in this mode, we must create a new zookeeper client
                    logger.info("the connection to zookeeper is expired!, retry...");
                    try {
                        logger.info("close the old zookeeper session...");
                        zkHelper.getZooKeeper().close();
                        logger.info("reconnecting to the zookeeper...");
                        zkHelper.init();
                    }
                    catch (Exception e){
                        logger.warn("error occured when reconnection to the zookeeper, {}", e.getMessage());
                    }
                }
                break;
            }

        }
    }

    private static class BrokerWatcher implements Watcher {

        private ZkHelper zkHelper;
        private List<BrokerInfo> brokerInfos = new LinkedList<>();

        public BrokerWatcher(ZkHelper zkHelper) {
            this.zkHelper = zkHelper;
        }

        @Override
        public void process(WatchedEvent event) {
            if (event.getType() == NodeChildrenChanged) {
                try {
                    logger.info("BrokerWatcher::process(...) is called. event.path: {}", event.getPath());
                    List<String> result = new LinkedList<>();
                    List<String> childrens = zkHelper.getZooKeeper().getChildren(event.getPath(), this);
                    childrens.forEach((s) -> {
                        byte[] data;
                        try {
                            String tmpPath = event.getPath() + "/" + s;
                            data = zkHelper.getZooKeeper().getData(tmpPath, false, null);
                            if (data != null) {
                                String v = new String(data);
                                result.add(v);
                            }
                        } catch (Exception e) {
                            logger.error("get data failed! path:{}, exception:{}", s, e.getMessage());
                        }
                    });
                    zkHelper.brokerMonitor.onBrokerChange(result);
                    // file the processor
                } catch (KeeperException | InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
    }
}
